//
// Created by fys on 2020/5/14.
//

#include "common_ttcp.h"

static int acceptOrDie(uint16_t port)
{
    int listenfd = socket(AF_INET, SOCK_STREAM, 0);
    assert(listenfd >= 0);

    int on = 1;
    if (setsockopt(listenfd, SOL_SOCKET, SO_REUSEADDR, &on, sizeof(on)))
    {
        perror("set socket failed!");
        exit(EXIT_FAILURE);
    }

    struct sockaddr_in addr;
    bzero(&addr, sizeof(addr));
    addr.sin_family = AF_INET;
    addr.sin_port = htons(port);
    addr.sin_addr.s_addr = htonl(INADDR_ANY);
    if (bind(listenfd, (struct sockaddr *)&addr, sizeof(addr)))
    {
        perror("bind failed!");
        exit(EXIT_FAILURE);
    }

    if (listen(listenfd, 1024))
    {
        perror("listen failed!");
        exit(EXIT_FAILURE);
    }

    socklen_t length;
    struct sockaddr_in clientaddr;
    length = sizeof(clientaddr);
    int connfd = accept(listenfd, (struct sockaddr *) &clientaddr , &length);
    if (connfd < 0)
    {
        perror("accept failed!");
        exit(EXIT_FAILURE);
    }
    close(listenfd);
    return connfd;
}

static struct sockaddr_in resolveOrDie(const char* host, uint16_t port)
{
    struct sockaddr_in serveraddr;
    bzero(&serveraddr, sizeof(serveraddr));
    serveraddr.sin_family = AF_INET;
    serveraddr.sin_port = htons(port);
    inet_pton(AF_INET, host, &serveraddr.sin_addr);
    return serveraddr;
}

static time_t getCurrentTimeStamp()
{
    struct timeb t;
    ftime(&t);
    return 1000 * t.time + t.millitm;
}

void receive(const Options* opt)
{
    printf("start receive server!\n");
    int connfd = acceptOrDie(opt->port);

    SessionMessage sessionMessage = {0, 0};
    if (readn(connfd, &sessionMessage, sizeof(sessionMessage)) != sizeof(sessionMessage))
    {
        perror("read SessionMessage failed!\n");
        exit(EXIT_FAILURE);
    }

    sessionMessage.number = ntohl(sessionMessage.number);
    sessionMessage.length = ntohl(sessionMessage.length);
    printf("receive number = %d\nreceive length = %d\n", sessionMessage.number, sessionMessage.length);

    const size_t total_len = sizeof(int32_t) + sessionMessage.length;
    PayloadMessage* payloadMessage = malloc(total_len);
    assert(payloadMessage);

    for (int i = 0; i < sessionMessage.number; i++)
    {
        printf("%d receive message!\n", i);
        // receive length
        payloadMessage->length = 0;
        if (readn(connfd, &payloadMessage->length, sizeof(payloadMessage->length)) != sizeof(payloadMessage->length))
        {
            perror("read length failed!");
            exit(EXIT_FAILURE);
        }
        payloadMessage->length = htonl(payloadMessage->length);
        assert(payloadMessage->length == sessionMessage.length);

        // receive date
        if (readn(connfd, payloadMessage->data, payloadMessage->length) != payloadMessage->length)
        {
            perror("read data failed!");
            exit(EXIT_FAILURE);
        }

        // return ack to client
        int32_t ack = htonl(payloadMessage->length);
        if (writen(connfd, &ack, sizeof(ack)) != sizeof(ack))
        {
            perror("return ack failed!");
            exit(EXIT_FAILURE);
        }
    }
    free(payloadMessage);
    close(connfd);
}

void transmit(const Options* opt)
{
    printf("start transmit client!\n");
    int socketfd = socket(AF_INET, SOCK_STREAM, 0);
    struct sockaddr_in addr = resolveOrDie(opt->host, opt->port);

    if (connect(socketfd, (struct sockaddr*) &addr, sizeof(addr)))
    {
        perror("connect failed!\n");
        close(socketfd);
        return;
    }
    printf("connect success!\n");

    // send SessionMessage include number and length
    SessionMessage sessionMessage = {htonl(opt->number), htonl(opt->length)};
    if (writen(socketfd, &sessionMessage, sizeof(sessionMessage)) != sizeof(sessionMessage))
    {
        perror("write session message failed!\n");
        exit(EXIT_FAILURE);
    }
    printf("write session message success!\n");

    // init send data
    PayloadMessage* payloadMessage = malloc(sizeof(PayloadMessage));
    assert(payloadMessage);
    payloadMessage->length = htonl(opt->length);
    for (int i = 0 ; i<opt->length ; i++)
    {
        payloadMessage->data[i] = "0123456789ABCDEF"[i % 16];
    }
    double total_mb = 1.0 * opt->length * opt->number / 1024 /1024;
    printf("%.3f MiB in total\n", total_mb);

    time_t startTime = getCurrentTimeStamp();
    size_t total_len = sizeof(int32_t) + opt->length;
    for (int i = 0 ; i < opt->number ; i++)
    {
        size_t write_num = writen(socketfd, payloadMessage, total_len);
        assert(write_num == total_len);

        int32_t ack = 0;
        size_t readnum = readn(socketfd, &ack, sizeof(int32_t));
        if (readnum == 0)
        {
            break;
        }
        assert(readnum == sizeof(int32_t));
        ack = ntohl(ack);
        assert(ack == opt->length);
    }

    free(payloadMessage);
    close(socketfd);
    time_t endTime = getCurrentTimeStamp();
    double diff =  ((double)(endTime - startTime))/1000;
    printf("%.3f seconds\n%.3f MiB/s\n", diff, total_mb / diff);
}